import numpy as np
import math
import datetime
import pandas as pd
import json
import gc


def complex_2_float(complex_num):
    if "j" in str(complex_num):
        complex_2_str = str(complex_num)
        try:
            if "+" in complex_2_str or complex_2_str.startswith("(-"):
                complex_num_temp = complex_2_str.split("+")[0]
                if "-" in complex_num_temp:
                    complex_num_temp = complex_num_temp.split("-")[0]
                    float_num = -float(complex_num_temp.replace("(", "0"))
                else:
                    float_num = float(complex_num_temp.replace("(", "0"))

            elif "-" in complex_2_str:
                float_num = float((complex_2_str.split("-")[0]).replace("(", "0"))
        except Exception as e:
            print("error --> {} {}".format(complex_2_str, complex_num))
            print(e)
    else:
        float_num = float(complex_num)
    return float_num


def product_dict_5(xs, xlist_o, coeffs, adding_ranges=5):
    """
    xs = 根的X坐标
    ys = 根的Y坐标
    coeffs = 参数
    """
    xlist = xlist_o.copy()
    p = np.poly1d(coeffs)
    a = 0
    xsL = xs.tolist()
    tx = []
    for _xss in xsL:
        _xs = complex_2_float(_xss)
        tx.append(_xs)

    xss = max(tx)
    tx1 = tx[1]
    tx0 = tx[0]

    for i in range(adding_ranges):
        xlist.append(i * 1.0)
    xlist.sort()

    newYhat = p(xlist)
    newdic = dict(zip(xlist, newYhat))

    return newdic


def product_dict(xs, ys, coeffs):
    """
    xs = 根的X坐标
    ys = 根的Y坐标
    coeffs = 参数
    """

    p = np.poly1d(coeffs)
    newX = []
    a = 0
    xsL = xs.tolist()
    ysL = ys.tolist()
    tx = []
    ty = []
    for _xss in xsL:
        _xs = complex_2_float(_xss)
        _ys = ysL[xsL.index(_xss)]
        _ys = complex_2_float(_ys)
        tx.append(_xs)
        ty.append(_xs)

    xss = max(tx)
    tx1 = tx[1]
    tx0 = tx[0]
    if tx1 < 0 and tx0 < 0:
        _tx = abs(abs(tx1) - abs(tx0))
    elif tx1 < 0 or tx0 < 0:
        _tx = abs(tx1) + abs(tx0)
    else:
        _tx = abs(tx1) - abs(tx0)

    for i in range(int(_tx)):
        newX.append(float(i - min(tx)))
    newYhat = p(newX)
    newdic = dict(zip(newX, newYhat))

    return newdic


def polyfit(x, y, degree):
    results = {}
    coeffs = np.polyfit(x, y, degree)
    results["polynomial"] = coeffs.tolist()

    # r-squared
    p = np.poly1d(coeffs)
    Q = np.polyder(coeffs)  # 求导数
    xs = np.roots(Q)  # 求多项式函数根
    ys = np.polyval(coeffs, xs)
    # fit values, and mean

    yhat = p(x)  # or [p(z) for z in x]
    ybar = np.sum(y) / len(y)  # or sum(y)/len(y)
    ssreg = np.sum((yhat - ybar) ** 2)  # or sum([ (yihat - ybar)**2 for yihat in yhat])
    sstot = np.sum((y - ybar) ** 2)  # or sum([ (yi - ybar)**2 for yi in y])

    if sstot != 0:
        results["determination"] = ssreg / sstot  # 准确率
    else:
        results["determination"] = 0
    results["Y"] = yhat  # TODO 要删除
    results["Q"] = Q
    results["xs"] = xs
    results["ys"] = ys
    return results


def polyfitlize(df, key, gap=5, max_incrace=0.2, pre="5"):
    """
    对数据进行拟合
    """
    l = []
    XYl = []
    YDic = {}
    determination_dic = {}
    XYdic = {}

    if gap > 40:
        ranges = gap * 2
    elif gap > 10:
        ranges = gap * 3
    else:
        ranges = gap * 4
    starttime = datetime.datetime.now()
    date_list = df.date.tolist()
    print("{} polyfitlize start".format(key), end=" ")

    total = (
        df.dropna(axis="index", how="all", subset=["high", "open"]).shape[0] - ranges
    )
    prescale = 0

    for _ in range(total):
        scale = int((_ / total) * 100)
        if scale % 5 == 0 and prescale != scale:
            print("*", end="")
        prescale = scale

        y = df.iloc[_ : ranges + _][key]
        y.dropna(axis=0, inplace=True)
        while np.nan in y:
            y.remove(np.nan)

        dic = {}
        dic["date"] = df.iloc[_]["date"]
        # x = [float(i * 1.0) for i in range(len(y))]
        x = []
        a = 0
        for i in y:
            a += 1
            x.append(float(a))

        results = polyfit(x, y, 3)
        # dic["determination"] = results["determination"]
        a = 0
        coeffsL = results["polynomial"]

        for i in coeffsL:
            dic["polynomial_X_{}".format(a)] = coeffsL[a]
            a += 1

        a = 0
        Ql = results["Q"]
        for i in Ql:
            dic["Q_{}".format(a)] = Ql[a]
            a += 1

        a = 0
        xsL = results["xs"].tolist()
        ysL = results["ys"].tolist()

        temp_dic = {}
        for xs in xsL:
            _xs = round(complex_2_float(xs) + _)
            if 0 < _xs < df.shape[0]:
                ys = ysL[xsL.index(xs)]
                ys = complex_2_float(ys)
                if int(_xs) > _:
                    day_gap = int(df.date.iloc[_ : int(_xs)].shape[0] / gap)
                elif int(_xs) == _:
                    day_gap = 1
                else:
                    day_gap = int(df.date.iloc[int(_xs) : _].shape[0] / gap)
                if (
                    df.iloc[_].Avg * ((1 - max_incrace) ** abs(day_gap))
                    <= ys
                    <= df.iloc[_].Avg * ((1 + max_incrace) ** abs(day_gap))
                    and day_gap < (240 / gap) * 5
                ):
                    temp_dic.update({df.iloc[int(_xs)].date: ys})

            if temp_dic != {}:
                XYdic[df.iloc[_].date] = temp_dic

        l.append(dic)
        YDic[df.iloc[_].date] = product_dict_5(
            results["xs"], x, results["polynomial"], 10
        )
        # T = dict(zip(x, results["Y"]))
        # YDic[df.iloc[_].date] = T

        determination_dic[df.iloc[_].date] = results["determination"]
    print(" polyfitlize => {}s".format((datetime.datetime.now() - starttime).seconds))

    if l != []:
        df2 = pd.DataFrame(l)
        df = pd.merge(df, df2, on="date", how="left")
    print(date_list[-1], len(date_list))

    print("Y_prodict   start", end=" ")
    new_dict = {}
    date_list = df.date.tolist()
    predic = {}
    for index in range(len(YDic.keys())):
        scale = int((index / len(YDic.keys())) * 100)
        if scale % 5 == 0 and prescale != scale:
            print("*", end="")
        prescale = scale
        date = date_list[index]
        value = YDic[date]
        if date in XYdic:
            if len(value) == 1:
                temp = [XYdic[date][i] for i in XYdic[date]][0]
                values_list = list(value.values())
                keys_list = list(value.keys())
                mx = abs(abs(max(values_list)) - temp)
                mi = abs(abs(min(values_list)) - temp)
                mxind = values_list.index(max(values_list))
                miind = values_list.index(min(values_list))
                cutter = int(len(values_list) / 2)

                if mx < mi:  # 更接近小数值根部
                    if mxind > miind:
                        keys_list = keys_list[cutter:mxind]
                        values_list = values_list[cutter:mxind]
                    else:
                        keys_list = keys_list[mxind:cutter]
                        values_list = values_list[mxind:cutter]

                elif mi < mx:  # 更接近大数值根部
                    if mxind > miind:
                        keys_list = keys_list[miind:cutter]
                        values_list = values_list[miind:cutter]
                    else:
                        keys_list = keys_list[cutter:miind]
                        values_list = values_list[cutter:miind]
                value = dict(zip(keys_list, values_list))
                # print(value)

            for ind in value:
                if 0 <= index + int(ind) <= len(date_list) and abs(ind) <= ranges:
                    if date_list[index + int(ind)] not in new_dict:
                        new_dict[date_list[index + int(ind)]] = {
                            "Y_prodict": round(value[int(ind)], 2),
                            "determination": determination_dic[date],
                            "prodict_date": date,
                        }
                    elif (
                        new_dict[date_list[index + int(ind)]]["determination"]
                        < determination_dic[date]
                    ):

                        new_dict[date_list[index + int(ind)]] = {
                            "Y_prodict": round(value[int(ind)], 2),
                            "determination": determination_dic[date],
                            "prodict_date": date,
                        }

    print("Y_prodict  => {}s".format((datetime.datetime.now() - starttime).seconds))

    _dfY = pd.DataFrame(new_dict).T.reset_index()
    _dfY.rename(columns={"index": "date"}, inplace=True)
    df = pd.merge(df, _dfY, on="date", how="left")

    new_l = []

    for XY in XYdic:
        dic = {}
        dic["date"] = XY
        if XYdic[XY] != {}:
            a = 0
            for _ in XYdic[XY]:
                dic["xss_{}".format(a)] = _
                dic["ys_{}".format(a)] = XYdic[XY][_]
                a += 1
            new_l.append(dic)

    XYdic = pd.DataFrame(new_l)
    # XYdic.rename(columns={"index": "date"}, inplace=True)
    df = pd.merge(df, XYdic, on="date", how="left")
    df.to_csv("Main.csv")
    return df